package org.zjt.agent.generate.javassist;

import javassist.*;
import lombok.extern.slf4j.Slf4j;
import org.zjt.agent.ClassDefineDescription;
import org.zjt.agent.Order;
import org.zjt.agent.generate.ClassHandler;
import org.zjt.agent.generate.GeneratorEnum;
import org.zjt.agent.generate.MethodRedefineStrategy;

import java.io.ByteArrayInputStream;
import java.util.*;
import java.util.function.Consumer;
import java.util.function.Predicate;

/**
 * Description:
 *
 * @author juntao.zhang
 * Date: 2018-07-06 下午5:44
 * @see
 */
@Slf4j
public class ThreadPoolClassHandler implements ClassHandler<ClassDefineDescription> {


    private List<MethodRedefineStrategy> methodRedefineStrategies = new ArrayList<>();



    public ThreadPoolClassHandler() {

    }

    @Override
    public Optional<byte[]> redefine(ClassDefineDescription classDefine) {
        try {
            CtClass ctClass = getCtClass(classDefine);
            handleMethods(ctClass);
            return Optional.of(ctClass.toBytecode());
        } catch (Exception e) {
            e.printStackTrace();
        }
        return Optional.empty();
    }

    @Override
    public boolean support(ClassDefineDescription classDefine) {
        boolean result = Objects.equals("java.util.concurrent.ThreadPoolExecutor", classDefine.getClassName()) ||
                Objects.equals("java.util.concurrent.ScheduledThreadPoolExecutor", classDefine.getClassName());
        return result;
    }

    @Override
    public int order() {
        return Order.HIGH.getValue();
    }

    @Override
    public GeneratorEnum getGenerator() {
        return GeneratorEnum.JAVASSIST;
    }

    @Override
    public void addMethodRedefineStrategies(List<MethodRedefineStrategy> methodRedefineStrategies) {
        this.methodRedefineStrategies.addAll(methodRedefineStrategies);
    }

    @Override
    public void addMethodRedefineStrategie(MethodRedefineStrategy methodRedefineStrategy) {
        methodRedefineStrategies.add(methodRedefineStrategy);
    }


    private void handleMethods(CtClass ctClass){
        CtMethod[] declaredMethods = ctClass.getDeclaredMethods();
        for (CtMethod ctMethod :declaredMethods){
            methodRedefineStrategies.stream().filter(methodRedefineStrategy -> methodRedefineStrategy.support(ctMethod)).forEach(methodRedefineStrategy -> methodRedefineStrategy.redefine(ctMethod));
        }
    }


    private static CtClass getCtClass(ClassDefineDescription classDefine)  {
       try {
           ClassLoader classLoader = classDefine.getClassLoader();
           byte[] classFileBuffer = classDefine.getClassFileBuffer();
           ClassPool classPool = new ClassPool(true);
           if (classLoader == null) {
               classPool.appendClassPath(new LoaderClassPath(ClassLoader.getSystemClassLoader()));
           } else {
               classPool.appendClassPath(new LoaderClassPath(classLoader));
           }
           CtClass clazz = classPool.makeClass(new ByteArrayInputStream(classFileBuffer), false);
           clazz.defrost();
           return clazz;
       }catch (Exception e){
           throw new RuntimeException(e);
       }
    }

}
